#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
* 说明：制作CSV格式
* 时间：2019-6-27

* 作者： whg
'''
import  xml.dom.minidom
import os,shutil,glob, csv
'''--------------------------配置文件------------------------'''
# XML文件的目录
xmlpath = '../src/xml'
# 图片的目录
photopath = '../src/photo'

xmlpath = os.path.join(os.getcwd(), xmlpath)
photopath = os.path.join(os.getcwd(), photopath)
def getxmlinfo(filename):
    dom = xml.dom.minidom.parse(filename)
    root = dom.documentElement
    namelist = root.getElementsByTagName('filename')
    name = namelist[0]
    label = list()
    objectlist = root.getElementsByTagName('object')
    for object in objectlist:
        flaglist = object.getElementsByTagName('name')
        flag = flaglist[0]
        xminlist = object.getElementsByTagName('xmin')
        xmin = xminlist[0]
        yminlist = object.getElementsByTagName('ymin')
        ymin = yminlist[0]
        xmaxlist = object.getElementsByTagName('xmax')
        xmax = xmaxlist[0]
        ymaxlist = object.getElementsByTagName('ymax')
        ymax = ymaxlist[0]
        label += list([flag.firstChild.data, xmin.firstChild.data, ymin.firstChild.data, xmax.firstChild.data, ymax.firstChild.data])
    return name.firstChild.data, label

class CSVDate():
    def __init__(self, photopath = 'photo', xmlpath = 'xml'):
        self.csv = []
        self.photopath = photopath
        self.xmlpath = xmlpath
        self.test_size = 0.2
        self.label = []

        self.getcsv()
        print(self.csv)
        self.write_file()

    def getcsv(self):
        xmllist = os.listdir(self.xmlpath)
        for xml in xmllist:
            _, date = getxmlinfo(os.path.join(self.xmlpath, xml))
            name = xml.split(".xml")[0]
            for i in range(int(len(date)/5)):
                if date[i*5 + 0][2] not in self.label:
                    self.label.append(date[i*5 + 0][2])
                self.csv.append(['data//src//photo//' + name + '.jpg', date[i*5 + 1], date[i*5 + 2], date[i*5 + 3], date[i*5 + 4],date[i*5 + 0][2]])

    def write_file(self):
        with open('train_annotations.csv', 'w', newline='') as fp:
            csv_writer = csv.writer(fp, dialect='excel')
            csv_writer.writerows(self.csv[int(len(self.csv)*self.test_size):])
            csv_writer.writerows(self.csv[:5])
        with open('val_annotations.csv', 'w', newline='') as fp:
            csv_writer = csv.writer(fp, dialect='excel')
            csv_writer.writerows(self.csv[:int(len(self.csv)*self.test_size)])

        class_name = sorted(self.label)
        class_ = []
        for num, name in enumerate(class_name):
            class_.append([name, num])
        with open('classes.csv', 'w', newline='') as fp:
            csv_writer = csv.writer(fp, dialect='excel')
            csv_writer.writerows(class_)


xml_file = glob.glob('./Annotations/*.xml')
CSVDate(photopath, xmlpath)

